#include "Pch.h"

#include "Dx11Shaders.h"
#include "Dx11Globals.h"
#include "Dx11Buffers.h"

#include "../../../MCD/Core/System/LinkList.h"
#include "../../../MCD/Core/System/Map.h"

#include <vector>

namespace MCD
{
	ID3DBlob* Dx11Shaders::compile(
		const char* srcData,
		size_t srcSize,
		const char* name,
		const char* entryPoint,
		const char* target
		)
	{
		ID3DBlob* code;
		ID3DBlob* errs;

		HRESULT hr = D3DCompile(
			srcData, srcSize, name,
			nullptr, nullptr,
			entryPoint, target,
			0, 0,
			&code, &errs);

		if(FAILED(hr))
		{
			std::wstring errStr = L"Unknown";
			if(errs)
				errStr = MCD::strToWStr((const char*)errs->GetBufferPointer());
			Log::format(Log::Error, (L"Dx11Shaders::compile failed; Reasons:\n" + errStr).c_str());
		}

		return code;
	}

	ID3D11VertexShader* Dx11Shaders::createVertexShader(ID3DBlob* byteCode)
	{
		MCD_ASSUME(nullptr != getD3D11Device());
		MCD_ASSUME(nullptr != byteCode);

		ID3D11VertexShader* shader;

		HRESULT hr = getD3D11Device()->CreateVertexShader(
			(DWORD*)byteCode->GetBufferPointer(), byteCode->GetBufferSize(),
			nullptr,
			&shader
			);

		if(FAILED(hr))
		{
			Log::format(Log::Error, L"Dx11Shaders::createVertexShader failed");
		}

		return shader;
	}

	ID3D11PixelShader* Dx11Shaders::createPixelShader(ID3DBlob* byteCode)
	{
		MCD_ASSUME(nullptr != getD3D11Device());
		MCD_ASSUME(nullptr != byteCode);

		ID3D11PixelShader* shader;

		HRESULT hr = getD3D11Device()->CreatePixelShader(
			byteCode->GetBufferPointer(), byteCode->GetBufferSize(),
			nullptr,
			&shader
			);

		if(FAILED(hr))
		{
			Log::format(Log::Error, L"Dx11Shaders::createPixelShader failed");
		}

		return shader;
	}

	const char* Dx11Shaders::getBestVertexShaderTarget()
	{
		switch(getD3D11Device()->GetFeatureLevel())
		{
		case D3D_FEATURE_LEVEL_11_0:
			return "vs_5_0";

		case D3D_FEATURE_LEVEL_10_1:
			return "vs_4_1";

		case D3D_FEATURE_LEVEL_10_0:
			return "vs_4_0";

		case D3D_FEATURE_LEVEL_9_3:
			return "vs_4_0_level_9_3";

		default:
			return "vs_4_0_level_9_1";
		}
	}

	const char* Dx11Shaders::getBestPixelShaderTarget()
	{
		switch(getD3D11Device()->GetFeatureLevel())
		{
		case D3D_FEATURE_LEVEL_11_0:
			return "ps_5_0";

		case D3D_FEATURE_LEVEL_10_1:
			return "ps_4_1";

		case D3D_FEATURE_LEVEL_10_0:
			return "ps_4_0";

		case D3D_FEATURE_LEVEL_9_3:
			return "ps_4_0_level_9_3";

		default:
			return "ps_4_0_level_9_1";
		}
	}

	/*------------------------------------------------------------------------*/
	class Dx11ShaderConstants::Impl
	{
	public:
		struct Constant : public MapBase<std::string>::Node<Constant>
		{
			size_t mOffset;
			size_t mSize;

			Constant(const char* name)
			{
				setKey(name);
			}
		};

		struct ConstantBuffer : public LinkListBase::Node<ConstantBuffer>
		{
			char* mBufPtr;
			size_t mBufSize;
			ID3D11Buffer* mDx11ConstBuf;

			Map<Constant> mConstants;

			ConstantBuffer(size_t bufSize, ID3D11Buffer* dx11ConstBuf)
			{
				mBufPtr = new char[bufSize];
				mBufSize = bufSize;
				memset(mBufPtr, 0, mBufSize);

				mDx11ConstBuf = dx11ConstBuf;
			}

			~ConstantBuffer()
			{
				delete[] mBufPtr;
				safeRelease(mDx11ConstBuf);
			}

			void commit()
			{
				D3D11_MAPPED_SUBRESOURCE mappedRes;

				HRESULT hr = getD3D11DeviceContextIM()->Map(mDx11ConstBuf, 0, D3D11_MAP_WRITE_DISCARD, 0, &mappedRes);

				if(FAILED(hr))
					return;

				memcpy(mappedRes.pData, mBufPtr, mBufSize);

				getD3D11DeviceContextIM()->Unmap(mDx11ConstBuf, 0);
			}
		};

		LinkList<ConstantBuffer> mConstBufs;
		std::vector<ID3D11Buffer*> mDx11ConstBufs;
		bool mDirty;
		size_t mCount;

		Impl() : mDirty(false), mCount(0)
		{
		}

		~Impl()
		{
			for(size_t i=0; i<mDx11ConstBufs.size(); ++i)
				safeRelease(mDx11ConstBufs[i]);
		}

		void commit()
		{
			if(!mDirty)
				return;

			for(ConstantBuffer* buf = mConstBufs.begin(); buf != mConstBufs.end(); buf = buf->next())
				buf->commit();
			
			mDirty = false;
		}
	};

	Dx11ShaderConstants::Dx11ShaderConstants(ID3DBlob* shaderByteCode)
		: mImpl(new Impl)
	{
		ID3D11ShaderReflection* shaderRefl;
		HRESULT hr = D3DReflect(shaderByteCode->GetBufferPointer(), shaderByteCode->GetBufferSize(), IID_ID3D11ShaderReflection, (void**)&shaderRefl);

		if(SUCCEEDED(hr))
		{
			D3D11_SHADER_DESC shaderdesc;
			shaderRefl->GetDesc(&shaderdesc);

			for(UINT i=0; i<shaderdesc.ConstantBuffers; ++i)
			{
				ID3D11ShaderReflectionConstantBuffer* constBufRefl = shaderRefl->GetConstantBufferByIndex(i);
				D3D11_SHADER_BUFFER_DESC constbufdesc;
				constBufRefl->GetDesc(&constbufdesc);

				// create dx11 constant buffer
				ID3D11Buffer* dx11buf = Dx11Buffers::createConstant(constbufdesc.Size);
				mImpl->mDx11ConstBufs.push_back(dx11buf);

				// create Impl::ConstantBuffer
				Impl::ConstantBuffer* constbuf = new Impl::ConstantBuffer(constbufdesc.Size, dx11buf);
				mImpl->mConstBufs.pushBack(*constbuf);

				for(UINT j=0; j<constbufdesc.Variables; ++j)
				{
					ID3D11ShaderReflectionVariable * constRefl = constBufRefl->GetVariableByIndex(j);
					D3D11_SHADER_VARIABLE_DESC constdesc;
					constRefl->GetDesc(&constdesc);
					
					// create Impl::Constant
					Impl::Constant* c = new Impl::Constant(constdesc.Name);
					c->mOffset = constdesc.StartOffset;
					c->mSize = constdesc.Size;

					constbuf->mConstants.insert(*c);
					MCD_VERIFY(c->mOffset + c->mSize <= constbuf->mBufSize);

					++mImpl->mCount;
				}

			}

			safeRelease(shaderRefl);
		}
	}

	Dx11ShaderConstants::~Dx11ShaderConstants()
	{
		delete mImpl;
	}

	size_t Dx11ShaderConstants::count() const
	{
		return mImpl->mCount;
	}

	void Dx11ShaderConstants::write(const char* name, void* dataPtr, size_t dataSize)
	{
		std::string name_str(name);

		for(Impl::ConstantBuffer* buf = mImpl->mConstBufs.begin(); buf != mImpl->mConstBufs.end(); buf = buf->next())
		{
			Impl::Constant* c = buf->mConstants.find(name_str);
			if(nullptr == c)
				continue;

			if(dataSize <= c->mSize)
				memcpy(buf->mBufPtr + c->mOffset, dataPtr, dataSize);

			// we find it
			mImpl->mDirty = true;
			break;
		}
	}

	void Dx11ShaderConstants::bindVS()
	{
		mImpl->commit();
		getD3D11DeviceContextIM()->VSSetConstantBuffers(0, mImpl->mDx11ConstBufs.size(), &mImpl->mDx11ConstBufs[0]);
	}

	void Dx11ShaderConstants::bindGS()
	{
		mImpl->commit();
		getD3D11DeviceContextIM()->GSSetConstantBuffers(0, mImpl->mDx11ConstBufs.size(), &mImpl->mDx11ConstBufs[0]);
	}

	void Dx11ShaderConstants::bindPS()
	{
		mImpl->commit();
		getD3D11DeviceContextIM()->PSSetConstantBuffers(0, mImpl->mDx11ConstBufs.size(), &mImpl->mDx11ConstBufs[0]);
	}

}	// namespace MCD